package lab.partitioner

import org.apache.spark.{Partitioner, SparkConf, SparkContext}

/**
  * Created by kawhi on 30/03/2017.
  * Edited by Q.
  * This is a test for PartialKey.scala
  */

object PartialKey {
  def main(args: Array[String]) {
    val conf = new SparkConf().setAppName("Partial Key Partitioner").setMaster("local[4]")
    val sc = new SparkContext(conf)
    val ports = sc.parallelize(List(
      "a1", 'b', 'c', 'a', 'a',
      'a', 'a', 'a', 'a', 'a',
      'a', 'a', 'a', 'a', 'a',
      'a', 'a', 'a', 'a', 'a',
      'a', 'a', 'a', 'a'), 4).map((_, 1))

    val wbt = ports.partitionBy(new PartialKeyPartitioner(3))
    /*.reduceByKey((a, b)=>a+b)*/
    val res = wbt.mapPartitionsWithIndex(showPartition).collect

    sc.stop()
  }

  def showPartition[T](i: Int, iter: Iterator[T]): Iterator[(T)] = {
    iter.foreach(ele => {
      println("partition\t" + i + "\t" + ele)
    })
    iter
  }
}


class PartialKeyPartitioner(partitions: Int) extends Partitioner {
  def numPartitions: scala.Int = {
    partitions
  }

  def getPartition(key: scala.Any): scala.Int = {
    val partition_1 = (key.hashCode() & Integer.MAX_VALUE) % partitions
    val partition_2 = ((key.hashCode() & Integer.MAX_VALUE) + 1) % partitions

    //    println("before sending:")
    //    println("key: " + key + " two choice: " + partition_1.toString + " or " + partition_2.toString)
    //    println(partitionBucket.get(0) + " " + partitionBucket.get(1) + " " + partitionBucket.get(2) + " ")

    if (partitionBucket.get(partition_1) < partitionBucket.get(partition_2)) {
      partitionBucket.plus(partition_1)
      //      println("after sending:")
      //      println(partitionBucket.get(0) + " " + partitionBucket.get(1) + " " + partitionBucket.get(2) + " ")
      //      println()
      return partition_1
    }
    else {
      partitionBucket.plus(partition_2)
      //      println(partitionBucket.get(0) + " " + partitionBucket.get(1) + " " + partitionBucket.get(2) + " ")
      //      println("after sending:")
      //      println(partitionBucket.get(0) + " " + partitionBucket.get(1) + " " + partitionBucket.get(2) + " ")
      //      println()
      return partition_2
    }

  }

  //val partitionBucket = new PartitionBucket
  object partitionBucket {
    var buckets = new Array[Int](partitions)

    def get(position: Int): Int = {
      buckets(position)
    }

    def plus(position: Int) {
      buckets(position) = buckets(position) + 1
    }
  }

}



